package com.jboltai.capability.vdb.search;

import com.alibaba.fastjson.JSONObject;
import com.jboltai.util.other.StringUtil;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.function.Consumer;

public class VDBConditions {

    /**
     * 输出的字段，默认所有字段
     */
    private List<String> select;

    /**
     * 要匹配的向量化文本
     */
    private List<String> matchText;

    /**
     * 向量字段参数：若索引类型为 HNSW，设置检索参数 efConstruction，指定需要访问的候选向量的数目。取值范围[1,32768]，默认为10。
     * 若索引类型为 IVF 系列，设置参数 nprobe ，指定所需查询的单位数量。取值范围[1,nlist]其中 nlist 在创建 Collection 时已设置。
     */
    private JSONObject vectorParams;

    /**
     * 搜索一致性策略
     */
    private ConsistencyLevel consistencyLevel = ConsistencyLevel.DEFAULT;


    private List<VDBFilter> filter = new ArrayList<>(2);


    /**
     * 取多少条数据
     */
    private Integer limit;

    /**
     * 从第几条开始取
     */
    private Integer offset;

    private Float textRate;
    private Float vectorRate;

    private JSONObject extra;


    public Float getTextRate() {
        return textRate;
    }

    public void setTextRate(Float textRate) {
        this.textRate = textRate;
    }

    public Float getVectorRate() {
        return vectorRate;
    }

    public void setVectorRate(Float vectorRate) {
        this.vectorRate = vectorRate;
    }

    public List<String> getSelect() {
        return select;
    }

    public VDBConditions select(List<String> select) {
        this.select = select;
        return this;
    }

    public List<String> getMatchText() {
        return matchText;
    }


    public JSONObject getVectorParams() {
        return vectorParams;
    }

    public <T> T getVectorParam(String key, Class<T> type) {
        return StringUtil.convert((vectorParams == null || vectorParams.containsKey(key) == false) ? null : vectorParams.get(key).toString(), type);
    }


    /**
     * 如果集合使用的hnsw索引类型，需要配置参数ef,指定搜索范围
     *
     * @param efConstruction
     * @return
     */
    public VDBConditions hnswIndexParam(Integer efConstruction) {
        if (vectorParams == null) {
            vectorParams = new JSONObject();
        }
        vectorParams.put("ef", efConstruction);
        return this;
    }

    /**
     * 如果集合使用的ivf索引类型，需要配置参数nprobe,指定要查询的单位数
     *
     * @param nprobe
     * @return
     */
    public VDBConditions ivfIndexParam(Integer nprobe) {
        if (vectorParams == null) {
            vectorParams = new JSONObject();
        }
        vectorParams.put("nprobe", nprobe);
        return this;
    }


    /**
     * 设置向量查询参数
     *
     * @param vectorParams
     * @return
     */
    public VDBConditions vectorParams(JSONObject vectorParams) {
        this.vectorParams = vectorParams;
        return this;
    }

    public List<VDBFilter> getFilter() {
        return filter;
    }

    public VDBConditions filter(List<VDBFilter> filter) {
        this.filter.addAll(filter);
        return this;
    }

    /**
     * 设置要返回的字段
     *
     * @param select
     * @return
     */
    public VDBConditions select(String... select) {
        this.select = Arrays.asList(select);
        return this;
    }

    /**
     * 要通过向量化匹配的文本
     *
     * @param text
     * @return
     */
    public VDBConditions matchText(String... text) {
        this.matchText = Arrays.asList(text);
        return this;
    }


    /**
     * 要通过向量化匹配的文本
     *
     * @param text
     * @return
     */
    public VDBConditions matchText(String text) {
        this.matchText = Arrays.asList(text);
        return this;
    }

    /**
     * 拼接 =条件
     *
     * @param name
     * @param value
     * @return
     */
    public VDBConditions eq(String name, String value) {
        this.filter.add(new VDBFilter(name, value, FilterOpt.EQ));
        return this;
    }

    /**
     * 拼接 =条件
     *
     * @param name
     * @param value
     * @return
     */
    public VDBConditions eq(String name, Integer value) {
        this.filter.add(new VDBFilter(name, value, FilterOpt.EQ));
        return this;
    }

    /**
     * 拼接 =条件
     *
     * @param name
     * @param value
     * @return
     */
    public VDBConditions eq(String name, Long value) {
        this.filter.add(new VDBFilter(name, value, FilterOpt.EQ));
        return this;
    }

    /**
     * 拼接 !=条件
     *
     * @param name
     * @param value
     * @return
     */
    public VDBConditions notEq(String name, String value) {
        this.filter.add(new VDBFilter(name, value, FilterOpt.NOT_EQ));
        return this;
    }

    /**
     * 拼接 !=条件
     *
     * @param name
     * @param value
     * @return
     */
    public VDBConditions notEq(String name, Integer value) {
        this.filter.add(new VDBFilter(name, value, FilterOpt.NOT_EQ));
        return this;
    }

    /**
     * 拼接 !=条件
     *
     * @param name
     * @param value
     * @return
     */
    public VDBConditions notEq(String name, Long value) {
        this.filter.add(new VDBFilter(name, value, FilterOpt.NOT_EQ));
        return this;
    }


    /**
     * 字段值在列表值的范围中
     *
     * @param name
     * @param values
     * @return
     */
    public VDBConditions in(String name, String... values) {
        this.filter.add(new VDBFilter(name, values, FilterOpt.IN));
        return this;
    }

    /**
     * 字段值不在列表值的范围中
     *
     * @param name
     * @param values
     * @return
     */
    public VDBConditions notIn(String name, String... values) {
        this.filter.add(new VDBFilter(name, values, FilterOpt.NOT_IN));
        return this;
    }

    /**
     * 拼接 >条件
     *
     * @param name
     * @param value
     * @return
     */
    public VDBConditions gt(String name, Integer value) {
        this.filter.add(new VDBFilter(name, value, FilterOpt.GT));
        return this;
    }


    /**
     * 拼接 >条件
     *
     * @param name
     * @param value
     * @return
     */
    public VDBConditions gt(String name, Long value) {
        this.filter.add(new VDBFilter(name, value, FilterOpt.GT));
        return this;
    }

    /**
     * 拼接 <条件
     *
     * @param name
     * @param value
     * @return
     */
    public VDBConditions lt(String name, Integer value) {
        this.filter.add(new VDBFilter(name, value, FilterOpt.LT));
        return this;
    }


    /**
     * 拼接 <条件
     *
     * @param name
     * @param value
     * @return
     */
    public VDBConditions lt(String name, Long value) {
        this.filter.add(new VDBFilter(name, value, FilterOpt.LT));
        return this;
    }

    /**
     * 拼接 >=条件
     *
     * @param name
     * @param value
     * @return
     */
    public VDBConditions ge(String name, Integer value) {
        this.filter.add(new VDBFilter(name, value, FilterOpt.GE));
        return this;
    }


    /**
     * 拼接 >=条件
     *
     * @param name
     * @param value
     * @return
     */
    public VDBConditions ge(String name, Long value) {
        this.filter.add(new VDBFilter(name, value, FilterOpt.GE));
        return this;
    }

    /**
     * 拼接 <=条件
     *
     * @param name
     * @param value
     * @return
     */
    public VDBConditions le(String name, Integer value) {
        this.filter.add(new VDBFilter(name, value, FilterOpt.LE));
        return this;
    }


    /**
     * 拼接 <=条件
     *
     * @param name
     * @param value
     * @return
     */
    public VDBConditions le(String name, Long value) {
        this.filter.add(new VDBFilter(name, value, FilterOpt.LE));
        return this;
    }

    public VDBConditions or() {
        this.filter.add(new VDBFilter(null, null, FilterOpt.OR));
        return this;
    }


    public VDBConditions and() {
        this.filter.add(new VDBFilter(null, null, FilterOpt.AND));
        return this;
    }


    /**
     * 拼接 and not 条件
     *
     * @return
     */
    public VDBConditions andNot() {
        this.filter.add(new VDBFilter(null, null, FilterOpt.AND_NOT));
        return this;
    }


    /**
     * 拼接 or not 条件
     *
     * @return
     */
    public VDBConditions orNot() {
        this.filter.add(new VDBFilter(null, null, FilterOpt.OR_NOT));
        return this;
    }

    /**
     * 括号包裹
     *
     * @param consumer
     * @return
     */
    public VDBConditions bracket(Consumer<VDBConditions> consumer) {

        this.filter.add(new VDBFilter(null, null, FilterOpt.BRACKET_LEFT));
        consumer.accept(this);
        this.filter.add(new VDBFilter(null, null, FilterOpt.BRACKET_RIGHT));
        return this;
    }


    public Integer getLimit() {
        return limit;
    }

    public VDBConditions limit(Integer limit) {
        this.limit = limit;
        return this;
    }

    public Integer getOffset() {
        return offset;
    }

    public VDBConditions offset(Integer offset) {
        this.offset = offset;
        return this;
    }

    public ConsistencyLevel getConsistencyLevel() {
        return consistencyLevel;
    }

    public VDBConditions consistencyLevel(ConsistencyLevel consistencyLevel) {
        this.consistencyLevel = consistencyLevel;
        return this;
    }

    public JSONObject getExtra() {
        return extra;
    }

    /**
     * 设置额外的查询参数，只对某些特定的vdb有效
     * @param extra
     * @return
     */
    public VDBConditions setExtra(JSONObject extra) {
        this.extra = extra;
        return this;
    }
}
